// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/functions/prestosql/function_base_test.h>
#include <pollux/parse/type_resolver.h>

namespace kumo::pollux::functions::test {

class SliceTestBase : public FunctionBaseTest {
 protected:
  static void SetUpTestCase() {
    parse::registerTypeResolver();
    memory::MemoryManager::testingSetInstance({});
  }

  virtual void testSlice(
      const std::string& expression,
      const std::vector<VectorPtr>& parameters,
      const ArrayVectorPtr& expectedArrayVector) {
    auto result = evaluate<ArrayVector>(expression, make_row_vector(parameters));
    ::kumo::pollux::test::assertEqualVectors(expectedArrayVector, result);
    EXPECT_FALSE(expectedArrayVector->hasOverlappingRanges());
  }

  void basicTestCases() {
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4, 5}});
      auto expectedArrayVector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      testSlice("slice(C0, 1, 4)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2}});
      auto expectedArrayVector = make_array_vector<int64_t>({{1, 2}});
      testSlice("slice(C0, 1, 4)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4, 5}});
      auto expectedArrayVector = make_array_vector<int64_t>({{3, 4}});
      testSlice("slice(C0, 3, 2)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{3, 4}});
      testSlice("slice(C0, 3, 3)", {array_vector}, expectedArrayVector);
    }
    // Negative start index.
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{2, 3, 4}});
      testSlice("slice(C0, -3, 3)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{2, 3, 4}});
      testSlice("slice(C0, -3, 5)", {array_vector}, expectedArrayVector);
    }
    // Negative length.
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      POLLUX_ASSERT_THROW(
          testSlice(
              "slice(C0, 1, -1)",
              {array_vector, array_vector, expectedArrayVector},
              expectedArrayVector),
          "The value of length argument of slice() function should not be negative");
    }
    // 0 start index.
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      POLLUX_ASSERT_THROW(
          testSlice(
              "slice(C0, 0, 1)",
              {array_vector, array_vector, expectedArrayVector},
              expectedArrayVector),
          "SQL array indices start at 1");
    }
    // 0 length.
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      testSlice("slice(C0, 1, 0)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      testSlice("slice(C0, -2, 0)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      testSlice("slice(C0, -2, 0)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      testSlice("slice(C0, -5, 5)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      testSlice("slice(C0, -6, 5)", {array_vector}, expectedArrayVector);
    }
    {
      auto array_vector = make_array_vector<int64_t>({{1, 2, 3, 4}});
      auto expectedArrayVector = make_array_vector<int64_t>({{}});
      testSlice("slice(C0, -6, 5)", {array_vector}, expectedArrayVector);
    }
    {
      // The implementation is zero-copy, so floating point number comparison
      // won't have issue here since numerical representation remains the same.
      auto array_vector = make_array_vector<double>({{2.3, 2.3, 2.2}});
      auto expectedArrayVector = make_array_vector<double>({{2.3, 2.2}});
      testSlice("slice(C0, 2, 3)", {array_vector}, expectedArrayVector);
    }
    // String array.
    {
      auto array_vector = make_array_vector<StringView>({{"a", "b", "c", "d"}});
      auto expectedArrayVector = make_array_vector<StringView>({{"b", "c"}});
      testSlice("slice(C0, 2, 2)", {array_vector}, expectedArrayVector);
    }
    // Out of bound start index.
    {
      auto array_vector = make_array_vector<StringView>({{"a", "b", "c", "d"}});
      auto expectedArrayVector = make_array_vector<StringView>({{}});
      testSlice("slice(C0, 5, 2)", {array_vector}, expectedArrayVector);
    }
    // Out of bound length.
    {
      auto array_vector = make_array_vector<StringView>({{"a", "b", "c", "d"}});
      auto expectedArrayVector = make_array_vector<StringView>({{"b", "c", "d"}});
      testSlice("slice(C0, 2, 5)", {array_vector}, expectedArrayVector);
    }
  }
};

} // namespace kumo::pollux::functions::test
